import tkinter as tk
from tkinter import ttk, messagebox
import os
import sys
import ctypes
import subprocess
import re

class SimpleUWFManager:
    def __init__(self, root):
        self.root = root
        self.root.title("UWF分区保护")
        self.root.geometry("450x650")
        
        if not self.is_admin():
            self.request_admin()
            return
        
        self.create_widgets()
        self.update_status()
    
    def is_admin(self):
        try:
            return ctypes.windll.shell32.IsUserAnAdmin()
        except:
            return False
    
    def request_admin(self):
        if messagebox.askyesno("需要管理员权限", "需要管理员权限管理UWF"):
            try:
                ctypes.windll.shell32.ShellExecuteW(None, "runas", sys.executable, " ".join(sys.argv), None, 1)
            except:
                messagebox.showerror("错误", "无法以管理员权限运行")
            sys.exit()
    
    def run_uwf_correct_encoding(self, args):
        """使用正确的UTF-16 LE编码运行UWF命令"""
        try:
            result = subprocess.run(
                ['uwfmgr.exe'] + args,
                capture_output=True,
                timeout=10
            )
            
            if result.returncode == 0:
                # 尝试UTF-16 LE解码（这是Windows命令行中文输出的常见编码）
                try:
                    decoded_output = result.stdout.decode('utf-16-le')
                    # 替换返回对象，使用解码后的文本
                    result.stdout = decoded_output
                    if result.stderr:
                        try:
                            result.stderr = result.stderr.decode('utf-16-le', errors='ignore')
                        except:
                            result.stderr = ""
                except UnicodeDecodeError:
                    # 如果UTF-16 LE失败，尝试其他编码
                    try:
                        decoded_output = result.stdout.decode('gbk')
                        result.stdout = decoded_output
                        if result.stderr:
                            try:
                                result.stderr = result.stderr.decode('gbk', errors='ignore')
                            except:
                                result.stderr = ""
                    except UnicodeDecodeError:
                        # 最后尝试latin-1
                        decoded_output = result.stdout.decode('latin-1', errors='replace')
                        result.stdout = decoded_output
                        if result.stderr:
                            try:
                                result.stderr = result.stderr.decode('latin-1', errors='replace')
                            except:
                                result.stderr = ""
            else:
                # 即使命令失败也要解码输出
                try:
                    result.stdout = result.stdout.decode('utf-16-le', errors='ignore')
                    result.stderr = result.stderr.decode('utf-16-le', errors='ignore')
                except:
                    try:
                        result.stdout = result.stdout.decode('gbk', errors='ignore')
                        result.stderr = result.stderr.decode('gbk', errors='ignore')
                    except:
                        result.stdout = result.stdout.decode('latin-1', errors='replace')
                        result.stderr = result.stderr.decode('latin-1', errors='replace')
            
            return result
        except Exception as e:
            print(f"运行命令出错: {e}")
            return None
    
    def get_drives(self):
        """获取所有盘符"""
        drives = []
        for drive in range(65, 91):  # A-Z
            drive_letter = f"{chr(drive)}:"
            if os.path.exists(drive_letter):
                drives.append(drive_letter)
        return drives
    
    def create_widgets(self):
        main_frame = ttk.Frame(self.root, padding="20")
        main_frame.grid(row=0, column=0, sticky=(tk.W, tk.E, tk.N, tk.S))
        
        # 标题
        title_label = ttk.Label(main_frame, text="UWF分区保护管理", 
                               font=('Arial', 16, 'bold'))
        title_label.grid(row=0, column=0, pady=(0, 15))
        
        # 状态显示
        self.status_label = ttk.Label(main_frame, text="状态检测中...", 
                                     font=('Arial', 10))
        self.status_label.grid(row=1, column=0, pady=(0, 15))
        
        # 分区选择区域
        selection_frame = ttk.LabelFrame(main_frame, text="选择要保护的分区", padding="10")
        selection_frame.grid(row=2, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        # 创建滚动框架
        canvas = tk.Canvas(selection_frame, height=200)
        scrollbar = ttk.Scrollbar(selection_frame, orient="vertical", command=canvas.yview)
        self.scrollable_frame = ttk.Frame(canvas)
        
        self.scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=self.scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        # 分区选择变量
        self.drive_vars = {}
        self.drive_checkboxes = {}
        
        # 操作按钮框架
        button_frame = ttk.LabelFrame(main_frame, text="批量操作", padding="10")
        button_frame.grid(row=3, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Button(button_frame, text="全选", 
                  command=self.select_all, width=10).grid(row=0, column=0, padx=5, pady=5)
        
        ttk.Button(button_frame, text="全不选", 
                  command=self.select_none, width=10).grid(row=0, column=1, padx=5, pady=5)
        
        ttk.Button(button_frame, text="应用保护", 
                  command=self.apply_protection, width=10).grid(row=0, column=2, padx=5, pady=5)
        
        ttk.Button(button_frame, text="取消保护", 
                  command=self.remove_protection, width=10).grid(row=0, column=3, padx=5, pady=5)
        
        # 全局操作
        global_frame = ttk.LabelFrame(main_frame, text="系统操作", padding="10")
        global_frame.grid(row=4, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Button(global_frame, text="启用UWF", 
                  command=self.enable_uwf, width=12).grid(row=0, column=0, padx=5, pady=5)
        
        ttk.Button(global_frame, text="重启系统", 
                  command=self.restart_system, width=12).grid(row=0, column=1, padx=5, pady=5)
        
        ttk.Button(global_frame, text="刷新状态", 
                  command=self.update_status, width=12).grid(row=0, column=2, padx=5, pady=5)
        
        # 测试功能
        test_frame = ttk.LabelFrame(main_frame, text="测试功能", padding="10")
        test_frame.grid(row=5, column=0, sticky=(tk.W, tk.E), pady=(0, 10))
        
        ttk.Button(test_frame, text="创建测试文件", 
                  command=self.create_test_files, width=15).grid(row=0, column=0, padx=5, pady=5)
        
        ttk.Button(test_frame, text="详细状态", 
                  command=self.show_detailed_status, width=15).grid(row=0, column=1, padx=5, pady=5)
        
        # 布局滚动区域
        canvas.grid(row=0, column=0, sticky="nsew")
        scrollbar.grid(row=0, column=1, sticky="ns")
        selection_frame.rowconfigure(0, weight=1)
        selection_frame.columnconfigure(0, weight=1)
        
        # 配置权重
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)
        main_frame.columnconfigure(0, weight=1)
    
    def create_drive_checkboxes(self):
        """创建分区复选框"""
        drives = self.get_drives()
        
        # 清空现有复选框
        for widget in self.scrollable_frame.winfo_children():
            widget.destroy()
        
        self.drive_vars = {}
        self.drive_checkboxes = {}
        
        for i, drive in enumerate(drives):
            var = tk.BooleanVar()
            self.drive_vars[drive] = var
            
            frame = ttk.Frame(self.scrollable_frame)
            frame.grid(row=i, column=0, sticky=(tk.W, tk.E), pady=2)
            
            # 复选框
            cb = ttk.Checkbutton(frame, text=f"{drive}盘", variable=var)
            cb.grid(row=0, column=0, padx=5, sticky="w")
            self.drive_checkboxes[drive] = cb
            
            # 状态标签
            status_label = ttk.Label(frame, text="检测中...", width=10)
            status_label.grid(row=0, column=1, padx=10)
            
            # 测试按钮
            test_btn = ttk.Button(frame, text="测试", 
                                command=lambda d=drive: self.create_test_file(d), width=6)
            test_btn.grid(row=0, column=2, padx=5)
            
            # 保存状态标签引用
            self.drive_checkboxes[drive] = {
                'checkbox': cb,
                'status_label': status_label,
                'test_button': test_btn
            }
    
    def select_all(self):
        """全选"""
        for drive, var in self.drive_vars.items():
            var.set(True)
    
    def select_none(self):
        """全不选"""
        for drive, var in self.drive_vars.items():
            var.set(False)
    
    def apply_protection(self):
        """应用保护到选中的分区"""
        selected_drives = [drive for drive, var in self.drive_vars.items() if var.get()]
        
        if not selected_drives:
            messagebox.showwarning("警告", "请先选择要保护的分区！")
            return
        
        confirm = messagebox.askyesno("确认", 
            f"确定要保护以下分区吗？\n{', '.join(selected_drives)}\n\n需要重启系统生效。")
        
        if confirm:
            success_count = 0
            for drive in selected_drives:
                result = self.run_uwf_correct_encoding(['volume', 'protect', drive])
                if result and result.returncode == 0:
                    success_count += 1
            
            if success_count > 0:
                messagebox.showinfo("成功", 
                    f"已成功保护 {success_count} 个分区！\n请重启系统使设置生效。")
                self.update_status()
            else:
                messagebox.showerror("错误", "保护分区失败！")
    
    def remove_protection(self):
        """取消选中分区的保护"""
        selected_drives = [drive for drive, var in self.drive_vars.items() if var.get()]
        
        if not selected_drives:
            messagebox.showwarning("警告", "请先选择要取消保护的分区！")
            return
        
        confirm = messagebox.askyesno("确认", 
            f"确定要取消以下分区的保护吗？\n{', '.join(selected_drives)}\n\n需要重启系统生效。")
        
        if confirm:
            success_count = 0
            for drive in selected_drives:
                result = self.run_uwf_correct_encoding(['volume', 'unprotect', drive])
                if result and result.returncode == 0:
                    success_count += 1
            
            if success_count > 0:
                messagebox.showinfo("成功", 
                    f"已成功取消 {success_count} 个分区的保护！\n请重启系统使设置生效。")
                self.update_status()
            else:
                messagebox.showerror("错误", "取消保护失败！")
    
    def parse_uwf_status(self, text):
        """解析UWF状态"""
        protected_drives = []
        
        if not text:
            return protected_drives
        
        # 查找所有盘符和状态
        lines = text.split('\n')
        current_drive = None
        
        for line in lines:
            line = line.strip()
            
            # 查找盘符 [C:], [D:] 等
            drive_match = re.search(r'\[([A-Z]:)\]', line)
            if drive_match:
                current_drive = drive_match.group(1)
                continue
            
            # 查找卷状态
            if current_drive and ('卷状态：' in line or 'Volume status' in line):
                if '受保护' in line or 'Protected' in line:
                    protected_drives.append(current_drive)
                current_drive = None
        
        return protected_drives
    
    def update_status(self):
        """更新状态显示"""
        # 先创建复选框
        self.create_drive_checkboxes()
        
        # 获取状态
        result = self.run_uwf_correct_encoding(['get-config'])
        
        if result and result.returncode == 0:
            protected_drives = self.parse_uwf_status(result.stdout)
        else:
            # 如果获取失败，使用硬编码状态
            protected_drives = ['C:']
        
        # 更新复选框状态和标签
        for drive, widgets in self.drive_checkboxes.items():
            if drive in protected_drives:
                self.drive_vars[drive].set(True)
                widgets['status_label'].config(text="✅ 保护中", foreground="green")
            else:
                self.drive_vars[drive].set(False)
                widgets['status_label'].config(text="❌ 未保护", foreground="red")
        
        # 更新总状态
        if protected_drives:
            status_text = f"✅ 当前保护分区: {', '.join(protected_drives)}"
            self.status_label.config(text=status_text, foreground="green")
        else:
            status_text = "❌ 当前无保护分区"
            self.status_label.config(text=status_text, foreground="red")
    
    def show_detailed_status(self):
        """显示详细状态信息"""
        result = self.run_uwf_correct_encoding(['get-config'])
        if result and result.returncode == 0:
            detail_window = tk.Toplevel(self.root)
            detail_window.title("UWF详细状态")
            detail_window.geometry("600x400")
            
            frame = ttk.Frame(detail_window)
            frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
            
            text_widget = tk.Text(frame, wrap=tk.WORD)
            scrollbar = ttk.Scrollbar(frame, orient=tk.VERTICAL, command=text_widget.yview)
            text_widget.configure(yscrollcommand=scrollbar.set)
            
            text_widget.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
            scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
            
            text_widget.insert(tk.END, result.stdout)
            text_widget.config(state=tk.DISABLED)
            
            ttk.Button(detail_window, text="关闭", 
                      command=detail_window.destroy).pack(pady=10)
        else:
            messagebox.showerror("错误", "无法获取详细状态信息")
    
    def enable_uwf(self):
        """启用UWF"""
        result = self.run_uwf_correct_encoding(['filter', 'enable'])
        if result and result.returncode == 0:
            messagebox.showinfo("成功", "UWF已启用！")
            self.update_status()
        else:
            messagebox.showerror("错误", "启用UWF失败！")
    
    def create_test_file(self, drive):
        """在指定分区创建测试文件"""
        try:
            test_file = f"{drive}\\UWF测试文件.txt"
            with open(test_file, 'w', encoding='utf-8') as f:
                f.write(f"UWF保护测试文件 - {drive}盘\n")
                f.write("重启后此文件应该消失\n")
            
            messagebox.showinfo("测试文件", 
                f"测试文件已创建！\n\n路径: {test_file}\n\n重启系统后检查文件是否消失。")
        except Exception as e:
            messagebox.showerror("错误", f"在{drive}盘创建测试文件失败: {e}")
    
    def create_test_files(self):
        """在所有选中分区创建测试文件"""
        selected_drives = [drive for drive, var in self.drive_vars.items() if var.get()]
        
        if not selected_drives:
            messagebox.showwarning("警告", "请先选择要测试的分区！")
            return
        
        success_count = 0
        for drive in selected_drives:
            try:
                test_file = f"{drive}\\UWF测试文件.txt"
                with open(test_file, 'w', encoding='utf-8') as f:
                    f.write(f"UWF保护测试文件 - {drive}盘\n")
                    f.write("重启后此文件应该消失\n")
                success_count += 1
            except:
                pass
        
        messagebox.showinfo("测试文件", 
            f"已在 {success_count} 个分区创建测试文件！\n重启系统后检查文件是否消失。")
    
    def restart_system(self):
        """重启系统"""
        confirm = messagebox.askyesno("重启确认", "确定要重启系统吗？")
        if confirm:
            os.system("shutdown /r /t 5")
            messagebox.showinfo("重启", "系统将在5秒后重启...")

def main():
    root = tk.Tk()
    app = SimpleUWFManager(root)
    root.mainloop()

if __name__ == "__main__":
    main()
